"""
    红利价值
    * 银行股：红利、ROE成长性、PB估值
    * 重资产股票 + 剔除微利股：综合估值因子
    * 增强：
"""
import pandas as pd
from backTestAnaly.bactestEngine import BackTestEngine
from backTestAnaly.dataLoader import DataLoader
from data_resource.data_bases import engine
from strategy.pb_roe import get_prev_financialdate, get_rebalance
from utilities.utilities_strategy import stockPool_raw
from utilities.utilities_statics import factor_standardize
from utilities.utilities_func import set_limit_type


class DivQuality(BackTestEngine):
    def __init__(self, price: DataLoader, rebalance_dates: list, financial_date: dict,
                 benchmark_data: pd.DataFrame = None,
                 cash=100000):
        """
            price: 前复权行情价格参数，columns=[trading, code, close, open, volume]
        """
        super().__init__(data=price, cash=cash, benchmark_data=benchmark_data)  # 注入交易日数据
        self.strategy_pool = None  # 每报告期截面股票池
        self.init_pool = stockPool_raw()  # 初始股票池, columns=[code, industry]
        self.rebalance_dates = rebalance_dates  # 调仓日期, 通过get_rebalance()函数获取
        self.financial_date = financial_date  # 报告期数组 {key: 交易日, value: 报告期}，若报告期为非交易日，则key为后一交易日
        # 设置因子数据, columns = [code, trading, signal]

        # 组织价值的变化
        self.total_portvalue = {}

        # 每期选股结果记录
        self.historical_picks = []

    def __next__(self):
        # 交易主体
        result = super().__next__()  # 获取父类迭代器返回的：下一数据帧索引

        # 处理未完成订单
        if (len(self.paramsPipe.stocks_unsettled) > 0) and (self.today not in self.rebalance_dates):
            _unsettled_stocs = self.paramsPipe.get_stocks_unsettled()
            _unsettled_stocs.sort_values(by='direction', ascending=True, inplace=True)
            _unsettled = []
            for i in range(len(_unsettled_stocs)):
                _code = _unsettled_stocs.iloc[i, 0]
                _size = _unsettled_stocs.iloc[i, 1]
                _price = _unsettled_stocs.iloc[i, 2]
                _direction = _unsettled_stocs.iloc[i, 3]
                _limit_type = set_limit_type(_code)

                # 下单
                if _direction == -1:
                    self.order.sell(_code, result, _limit_type, _price, _size)
                    if self.order.order_status != 1:
                        _unsettled.append((_code, _size, _price, _direction))

                elif _direction == 1:
                    self.order.buy(_code, result, _limit_type, _price, _size)
                    if self.order.order_status != 1:
                        _unsettled.append((_code, _size, _price, _direction))
            # 对未完成订单进行替换
            self.paramsPipe.update_stocks_unsettled(_unsettled)

        """策略主体, 直接买入行业财务优选股票"""
        if (self.today in self.rebalance_dates) and (self.strategy_pool is not None):
            # self.paramsPipe.update_stocks_unsettled([])
            _f = pd.DataFrame(data={
                'code': self.strategy_pool,
                'target_weight': [0.85 / len(self.strategy_pool)] * len(self.strategy_pool)
            })
            _f['limit_type'] = _f['code'].apply(lambda c: set_limit_type(c))
            self.order.bulk_order_targetWeight(_f, result)
            self.strategy_pool = None  # 清空持仓

    def after_trading(self):
        today = self.today.strftime("%Y-%m-%d")
        if today in self.financial_date.keys():
            enddate = self.financial_date[today]  # 格式：%Y-%m-%d

            sql1 = f"""
                select a.code, a.industry_name, a.end_date, a.roe_q, a.roe_yoy_growth as roe_growth,
                    b.dv_ratio, c.bp, c.bp_pivot_1y, d.ep, d.ep_pivot_1y
                from quant_research.factor_finance_indicator as a
                left join quant_research.indicator_daily as b
                    on a.code=b.code and b.trade_date='{today}'
                left join quant_research.factor_bp_compound as c
                    on a.code=c.code and c.trade_date='{today}'
                left join quant_research.factor_ep_compound as d
                    on a.code=d.code and d.trade_date='{today}'
                where a.end_date='{enddate}'
            """
            _raw = pd.read_sql(sql1, engine)
            if _raw.empty:
                return
            else:
                # 银行股
                _raw_bank = _raw[_raw['industry_name'] == '银行'].copy()
                _raw_bank['dv_ratio'] = factor_standardize(_raw_bank['dv_ratio'], method='zscore', extreme=True)
                _raw_bank['roe_growth'] = factor_standardize(_raw_bank['roe_growth'], method='zscore', extreme=True)
                _raw_bank['roe_q'] = factor_standardize(_raw_bank['roe_q'], method='zscore', extreme=True)
                _raw_bank['bp'] = factor_standardize(_raw_bank['bp'], method='zscore', extreme=True)
                _raw_bank['bp_pivot_1y'] = factor_standardize(_raw_bank['bp_pivot_1y'], method='zscore', extreme=True)

                _raw_bank['signal'] = (_raw_bank['dv_ratio'] + _raw_bank['roe_growth'] + _raw_bank['bp'] +
                                       _raw_bank['bp_pivot_1y'] + _raw_bank['roe_q'])
                _raw_bank.sort_values('signal', ascending=False, inplace=True)
                _bank = _raw_bank['code'].tolist()[:2]

                # 非银行股
                _raw_non_bank = _raw[_raw['industry_name'] != '银行'].copy()
                _raw_non_bank['dv_ratio'] = factor_standardize(_raw_non_bank['dv_ratio'], method='zscore', extreme=True)
                _raw_non_bank['roe_growth'] = factor_standardize(_raw_non_bank['roe_growth'], method='zscore', extreme=True)
                _raw_non_bank['roe_q'] = factor_standardize(_raw_non_bank['roe_q'], method='zscore', extreme=True)
                _raw_non_bank['ep'] = factor_standardize(_raw_non_bank['ep'], method='zscore', extreme=True)
                _raw_non_bank['ep_pivot_1y'] = factor_standardize(_raw_non_bank['ep_pivot_1y'], method='zscore', extreme=True)
                _raw_non_bank['bp'] = factor_standardize(_raw_non_bank['bp'], method='zscore', extreme=True)
                _raw_non_bank['bp_pivot_1y'] = factor_standardize(_raw_non_bank['bp_pivot_1y'], method='zscore', extreme=True)

                _raw_non_bank['signal'] = (_raw_non_bank['dv_ratio'] + _raw_non_bank['roe_growth'] + _raw_non_bank['roe_q'] +
                                           _raw_non_bank['ep'] + _raw_non_bank['ep_pivot_1y'] + _raw_non_bank['bp'] +
                                           _raw_non_bank['bp_pivot_1y'])
                _raw_non_bank.sort_values('signal', ascending=False, inplace=True)
                _non_bank = _raw_non_bank['code'].tolist()[:28]

                # 股票池
                self.strategy_pool = _bank + _non_bank

                _f = pd.DataFrame(data={
                    'code': self.strategy_pool,
                    'target_weight': [0.85 / len(self.strategy_pool)] * len(self.strategy_pool)
                })
                self.historical_picks.append({self.today: _f})


def strategy_run():
    # 获取行情数据 2024-01-01至今，中证800+中证1000成分股
    sql1 = """
        select trade_date as trading, code, close, open, vol as volume, high, low
        from quant_research.market_daily_qfq
        where quant_research.market_daily_qfq.code in (
            -- 用中证1000, 中证800成分股做为股票池
            select con_code from quant_research.index_constituent where trade_date=(
                select max(trade_date) from quant_research.index_constituent
            ) and index_code in ('000852.SH', '000906.SH')
        ) and trade_date>='2014-08-31' and trade_date < '2025-08-30'
        order by trade_date;
    """
    # 获取bench_mark信息
    sql2 = """
        select trade_date as trading, close
        from quant_research.index_daily
        where ts_code='000300.SH' and trade_date >= '2014-08-31' and trade_date < '2025-08-30'
        order by trade_date;
    """
    market_data = pd.read_sql(sql1, engine)
    bench_mark = pd.read_sql(sql2, engine)
    # 准备数据管线
    price = DataLoader(market_data=market_data)
    # 报告期数组
    financial_dates_dict = get_prev_financialdate(2014, 2025)
    # 补充2025年
    financial_dates_dict['2025-03-31'] = '2025-03-31'
    financial_dates_form = []
    # 获取调仓日期
    for i in range(2014, 2025):
        financial_dates_form.append(f'{i}-04-30')
        financial_dates_form.append(f'{i}-08-31')
        financial_dates_form.append(f'{i}-10-31')

    rebalances = get_rebalance(financial_dates_form)
    # 注入回测模型
    _strategy = DivQuality(price=price, rebalance_dates=rebalances, financial_date=financial_dates_dict,
                           cash=100000000, benchmark_data=bench_mark)
    _strategy.run()
    return _strategy


def main():
    # 运行策略
    strategy = strategy_run()
    # 获取回测结果
    print("---------------- 1.组合分析 -----------------")
    strategy.analysis.port_analysis()
    print(strategy.analysis.port_indicators)
    strategy.analysis.plot_return(save_path='dividend_quality')
    return strategy


if __name__ == '__main__':
    main()
